import os
import sys
import numpy as np
import cv2

IMAGE_SIZE = 64


# 按照指定大小调整尺寸
def resize_image(image, height=IMAGE_SIZE, width=IMAGE_SIZE):
    top, bottom, left, right = (0, 0, 0, 0)

    # 获取图像尺寸
    h, w, _ = image.shape

    # 长宽不同的图片，找最大
    longest_edge = max(h, w)

    # 计算补全长宽的像素
    if h < longest_edge:
        dh = longest_edge - h
        top = dh // 2
        bottom = dh - top
    elif w < longest_edge:
        dw = longest_edge - w
        top = dw // 2
        right = dw - left
    else:
        pass

    # RGB
    Black = [0, 0, 0]

    # 补全图片
    constant = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                  value=Black)

    # 调整大小并返回
    return cv2.resize(constant, (height, width))


# 读取训练集数据
images = []
labels = []


def read_path(path_name):
    for dir_item in os.listdir(path_name):
        # 从初始路径开始叠加， 合并成可识别的操作路径
        full_path = os.path.abspath(os.path.join(path_name, dir_item))

        if os.path.isdir(full_path):
            read_path(full_path)
        else:
            if dir_item.endswith('.jpg'):
                image = cv2.imread(full_path)
                image = resize_image(image, IMAGE_SIZE, IMAGE_SIZE)

                # 该语句可以看到实际调用效果
                cv2.imwrite('1.jpg', image)

                images.append(image)
                labels.append(path_name)
    return images, labels


# 从指定路径读取训练集
def load_dataset(path_name):
    images, labels = read_path(path_name)

    # 将所有数据转化成四维数组， 尺寸为(图片数量*IMAGE_SIZE*IMAGE_SIZE*3)
    images = np.array(images)
    print(images.shape)

    # 标注数据
    labels = np.array([0 if label.endswith('duanchen') else 1 for label in labels])

    return images, labels


if __name__ == '__main__':
    if len(sys.argv) != 1:
        print("Usage:%s path_name\r\n" % (sys.argv[0]))
    else:
        images, labels = load_dataset("C:\\Users\\duanchen\\Documents\\Tencent Files\\1832931759\\FileRecv\\images\\duanchen")
